//! 连通性问题
//! ## UF 使用
//!
//! ```
//!     use algorithms_fourth::graph::connect::UF;
//!     let mut uf = UF::new(8);
//!     assert_eq!(uf.count(),8);
//!     uf.union(0, 1);
//!     assert_eq!(uf.count(),7);
//!     uf.union(2,3);
//!     assert_eq!(uf.count(),6);
//!     uf.union(4, 5);
//!     assert_eq!(uf.count(),5);
//!     uf.union(6, 7);
//!     assert_eq!(uf.count(),4);
//!     assert!(uf.connected(0, 1));
//!     assert!(!uf.connected(2, 1));
//!     uf.union(0, 2);
//!     uf.union(4,6);
//!     uf.union(0, 4);
//!     assert_eq!(uf.count(),1);
//! ```
//!
//! ## Connect 使用
//!
//! ```
//!    use algorithms_fourth::{
//!        add_edge,
//!        graph::{connect::Connect, Graph},
//!    };
//!
//!        let mut g = Graph::new(13);
//!        add_edge!(g, 0, 6, 5, 2, 1);
//!        g.add_edge(1, 0);
//!        g.add_edge(2, 0);
//!        add_edge!(g, 3, 5, 4);
//!        add_edge!(g, 4, 5, 6, 3);
//!        add_edge!(g, 5, 3, 4, 0);
//!        add_edge!(g, 6, 0, 4);
//!        add_edge!(g, 7, 8);
//!        add_edge!(g, 8, 7);
//!        add_edge!(g, 9, 11, 10, 12);
//!        add_edge!(g, 10, 9);
//!        add_edge!(g, 11, 9, 12);
//!        add_edge!(g, 12, 11, 9);
//!        // g.print_dot("graph.dot");
//!        let connect = Connect::new(&g);
//!        // connect.print_dot("connect.dot");
//!        assert_eq!(connect.count(), 3);
//!        assert_eq!(connect.connected(8, 7), true);
//!        assert_eq!(connect.connected(9, 10), true);
//!        assert_eq!(connect.connected(0, 6), true);
//!        assert_eq!(connect.connected(0, 8), false);
//!        assert_eq!(connect.connected(10, 12), true);
//!        assert_eq!(connect.id(6), 0);
//!        assert_eq!(connect.id(8), 1);
//!        assert_eq!(connect.id(12), 2);
//! ```
use std::fs;

use super::Graph;
pub struct UF {
    /// 索引
    id: Vec<usize>,
    /// 权重
    sz: Vec<usize>,
    /// 连通分量的数量
    count: usize,
}
impl UF {
    pub fn new(n: usize) -> Self {
        let mut id = vec![0; n];
        let sz = vec![1; n];
        (0..n).for_each(|i| {
            id[i] = i;
        });
        UF { id, sz, count: n }
    }
    pub fn count(&self) -> usize {
        self.count
    }
    pub fn connected(&self, p: usize, q: usize) -> bool {
        self.find(p) == self.find(q)
    }

    pub fn find(&self, p: usize) -> usize {
        let mut p = p;
        while p != self.id[p] {
            p = self.id[p];
        }
        p
    }
    pub fn union(&mut self, p: usize, q: usize) {
        let i = self.find(p);
        let j = self.find(q);
        if i != j {
            if self.sz[i] < self.sz[j] {
                self.id[i] = j;
                self.sz[j] += self.sz[i];
            } else {
                self.id[j] = i;
                self.sz[i] += self.sz[j];
            }
            self.count -= 1;
        }
    }
}
pub struct Connect {
    marked: Vec<bool>,
    id: Vec<usize>,
    count: usize,
}
impl Connect {
    pub fn new(g: &Graph) -> Self {
        let v = g.vertex_count();
        let mut marked = Vec::with_capacity(v);
        let mut id = Vec::with_capacity(v);
        for _ in 0..v {
            marked.push(false);
            id.push(usize::MAX);
        }
        let mut connect = Connect {
            marked,
            id,
            count: 0,
        };
        for s in 0..v {
            if !connect.marked[s] {
                connect.dfs(g, s);
                connect.count += 1;
            }
        }
        connect
    }

    fn dfs(&mut self, g: &Graph, v: usize) {
        self.marked[v] = true;
        self.id[v] = self.count;
        for w in g.adj(v).clone() {
            if !self.marked[w] {
                self.dfs(g, w);
            }
        }
    }
    pub fn connected(&self, v: usize, w: usize) -> bool {
        self.id[v] == self.id[w]
    }
    pub fn id(&self, v: usize) -> usize {
        self.id[v]
    }
    pub fn count(&self) -> usize {
        self.count
    }
    /// 打印dot格式的图结构
    pub fn print_dot(&self, path: &str) {
        let mut graph = r#"graph{
            label="graph"
             bgcolor=grey
             shape=circle
             "#
        .to_string();
        for t in 0..self.count {
            graph.push_str(&format!(
                "trees_{}[color=red;style=filled;label=\"{}\"]\n",
                t, t
            ));
        }
        for (v, w) in self.id.iter().enumerate() {
            graph.push_str(&format!("{} -- trees_{}\n", v, w));
        }
        graph.push('}');
        let _ = fs::write(path, graph);
    }
}

#[cfg(test)]
mod test {
    use crate::{
        add_edge,
        graph::{connect::Connect, Graph},
    };

    use super::UF;

    #[test]
    fn test() {
        let mut g = Graph::new(13);
        add_edge!(g, 0, 6, 5, 2, 1);
        g.add_edge(1, 0);
        g.add_edge(2, 0);
        add_edge!(g, 3, 5, 4);
        add_edge!(g, 4, 5, 6, 3);
        add_edge!(g, 5, 3, 4, 0);
        add_edge!(g, 6, 0, 4);
        add_edge!(g, 7, 8);
        add_edge!(g, 8, 7);
        add_edge!(g, 9, 11, 10, 12);
        add_edge!(g, 10, 9);
        add_edge!(g, 11, 9, 12);
        add_edge!(g, 12, 11, 9);
        // g.print_dot("graph.dot");
        let connect = Connect::new(&g);
        // connect.print_dot("connect.dot");
        assert_eq!(connect.count(), 3);
        assert_eq!(connect.connected(8, 7), true);
        assert_eq!(connect.connected(9, 10), true);
        assert_eq!(connect.connected(0, 6), true);
        assert_eq!(connect.connected(0, 8), false);
        assert_eq!(connect.connected(10, 12), true);
        assert_eq!(connect.id(6), 0);
        assert_eq!(connect.id(8), 1);
        assert_eq!(connect.id(12), 2);
    }
    #[test]
    fn test_union_find() {
        let mut uf = UF::new(8);
        assert_eq!(uf.count(), 8);
        uf.union(0, 1);
        assert_eq!(uf.count(), 7);
        uf.union(2, 3);
        assert_eq!(uf.count(), 6);
        uf.union(4, 5);
        assert_eq!(uf.count(), 5);
        uf.union(6, 7);
        assert_eq!(uf.count(), 4);
        assert!(uf.connected(0, 1));
        assert!(!uf.connected(2, 1));
        uf.union(0, 2);
        uf.union(4, 6);
        uf.union(0, 4);
        assert_eq!(uf.count(), 1);
    }
}
